# -*- coding: utf-8 -*-
# @Time    : 2024/7/30 17:15
# @Author  : yujiahao
# @File    : 71_fastapi_websocket_test.py
# @description: 测试 WebSockets

# 导入所需的模块
from fastapi import FastAPI
from fastapi.testclient import TestClient
from fastapi.websockets import WebSocket

# 创建 FastAPI 应用实例
app = FastAPI()


# 定义一个 HTTP GET 路由
@app.get("/")
async def read_main():
    """
    处理根路径的 GET 请求，返回一个简单的 JSON 响应。
    """
    return {"msg": "Hello World"}


# 定义一个 WebSocket 路由
@app.websocket("/ws")
async def websocket(websocket: WebSocket):
    """
    处理 WebSocket 连接，接受连接后发送一条 JSON 消息，然后关闭连接。
    """
    await websocket.accept()  # 接受 WebSocket 连接
    await websocket.send_json({"msg": "Hello WebSocket"})  # 发送 JSON 消息
    await websocket.close()  # 关闭 WebSocket 连接


# 测试 HTTP GET 请求的函数
def test_read_main():
    """
    测试根路径的 GET 请求，验证返回的状态码和 JSON 响应。

    测试 HTTP GET 请求：

    1. 创建一个 TestClient 实例。
    2. 使用客户端发送一个 GET 请求到 / 路由。
    3. 断言响应状态码为 200。
    4. 断言响应 JSON 数据为 {"msg": "Hello World"}。
    """
    client = TestClient(app)  # 创建一个 TestClient 实例
    response = client.get("/")  # 发送 GET 请求到根路径
    assert response.status_code == 200  # 断言响应状态码为 200
    assert response.json() == {"msg": "Hello World"}  # 断言响应的 JSON 数据


# 测试 WebSocket 连接的函数
def test_websocket():
    """
    测试 WebSocket 连接，验证接收到的 JSON 消息。

    测试 WebSocket 连接：

    1. 创建一个 TestClient 实例。
    2. 使用 with 语句连接到 /ws WebSocket 路由。
    3. 接收 WebSocket 发送的 JSON 数据。
    4. 断言接收到的数据为 {"msg": "Hello WebSocket"}。

    """
    client = TestClient(app)  # 创建一个 TestClient 实例
    # 使用 with 语句连接到 WebSocket 路由
    with client.websocket_connect("/ws") as websocket:
        data = websocket.receive_json()  # 接收 WebSocket 发送的 JSON 数据
        assert data == {"msg": "Hello WebSocket"}  # 断言接收到的数据
